package training

import (
	"MNIST-CNN/pkg/network"
	"fmt"
	"time"
)

// TrainCNN 训练CNN模型
func TrainCNN(nn *network.NeuralNetwork, trainDataset *Dataset, testDataset *Dataset, batchSize int, learningRate float64, epochs int, numClasses int) {
	// 获取图像维度（假设MNIST为28x28）
	imageHeight := 28
	imageWidth := 28

	// 准备训练数据
	trainInputs := network.PrepareImageBatch(trainDataset.Images, imageHeight, imageWidth)
	trainTargets := network.PrepareTargets(trainDataset.Labels, numClasses)

	// 准备测试数据
	testInputs := network.PrepareImageBatch(testDataset.Images, imageHeight, imageWidth)
	testTargets := network.PrepareTargets(testDataset.Labels, numClasses)

	// 训练前评估
	startEvaluation := time.Now()
	initialAccuracy := nn.Evaluate(testInputs, testTargets)
	evaluationTime := time.Since(startEvaluation)

	fmt.Printf("训练前 - 准确率: %.2f%%\n", initialAccuracy*100)
	fmt.Printf("评估耗时: %v\n", evaluationTime)

	// 训练模型
	startTrain := time.Now()
	nn.Train(trainInputs, trainTargets, batchSize, learningRate, epochs)
	trainTime := time.Since(startTrain)
	fmt.Printf("训练耗时: %v\n", trainTime)

	// 训练后评估
	finalAccuracy := nn.Evaluate(testInputs, testTargets)
	fmt.Printf("训练后 - 准确率: %.2f%%\n", finalAccuracy*100)
}
